clear
% All joints are revolute in this version
% The first body is attached to ground, the last body is free

% Initialize stuff
% Fixed parameters for the system. Each body in the system is identical
NBODIES = 16;
MASS = 10;
INERTIA = 10;
NP = 8; % Number of processors. Must be a power of 2 and less than NBODIES
NB_SYS = NBODIES/NP; % Number of bodies per subsystem at leaf level

% Center of mass located at the center of each link.
LENGTH = 1;
q(1:NBODIES,1) = 1;
u(1:NBODIES,1) = 1;
udot(1:NBODIES,1)=0;
Fh1 = zeros(6,NBODIES-1);
Fh2 = zeros(6,NBODIES-1);
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% This is just sorting of various indices. The only variables of interest
% here are SID and ASMID_T
% SID stores indices of the bodies in each subsystem.
    % For 12 body problem with 4 processors
    % SID =
    % 1     2     3 -> Processor 1 (assemblies 13 and 14)
    % 4     5     6 -> Processor 2 (assemblies 15 and 16)
    % 7     8     9 -> Processor 3 (assemblies 17 and 18)
    % 10    11    12 ->Processor 4 (assemblies 19 and 20)
% For a system with NBODIES, there are NBODIES-1 no of assemblies. ASMID_T 
% has four columns, first column is the index of hte parent body/assembly, 
% second of child body/assembly and third of the assembled parent-child. 
% Fourth column gives the index of udot associated with the assembly. For 
% a 12 body problem with 4 processors
% ASMID_T = 
%      1     2    13     2
%     13     3    14     3
%      4     5    15     5
%     15     6    16     6
%      7     8    17     8
%     17     9    18     9
%     10    11    19    11
%     19    12    20    12
%     14    16    21     4
%     18    20    22    10
%     21    22    23     7
% The part which sorts the indices is not really commented. It is not
% central to anything in the rest of the code. As long as components of
% ASMID_T are clear, you can go straight to line 100 and look at the rest
% of the code.
count_T = 1;
for i=1:NP    
    SID(i,:) = [NB_SYS*(i-1)+1:NB_SYS*i];
    ASMID(i,:) = [NBODIES+1+(NB_SYS-1)*(i-1):NBODIES+(NB_SYS-1)*i];
    
    prev = SID(i,1);
    for j=2:NB_SYS
        ASMID_T(count_T,:) = [prev,SID(i,j),ASMID(i,j-1),SID(i,j)];
        prev = ASMID(i,j-1);
        count_T = count_T + 1;
    end
end
PID = NBODIES+(NB_SYS-1)*NP+1:NBODIES+NB_SYS*NP-1;

for i=1:NP/2
    ASM_DISS(i,:)=[ASMID(2*i-1,end),ASMID(2*i,end)];
end

count = NP/2;
tmp_var = NP/2+1;
id = 1;
while(count>1)          
    for i=id:id+count/2-1
        ASM_DISS(tmp_var,:) = [PID(2*i-1),PID(2*i)];
        tmp_var = tmp_var + 1;
    end
    count = count/2;
    id = id+count/2+1;
end

for i=1:NP-1
    a = ASM_DISS(i,1);
    b = ASM_DISS(i,2);
    if(a > NBODIES && b > NBODIES) 
        while(a > NBODIES)            
            tmp_var_2 = find(~(ASMID_T(:,3)-a));
            a = ASMID_T(tmp_var_2,2);
        end
        tmp_var_3 = a+1;
    else
        tmp_var_1 = ASM_DISS(i,1);
        tmp_var_2 = find(~(ASMID_T(:,3)-tmp_var_1));
        tmp_var_3 = ASMID_T(tmp_var_2,4)+1;
    end
    ASMID_T(count_T,:) = [ASM_DISS(i,:),PID(i),tmp_var_3];
    count_T = count_T + 1;
end
% Sorting of indices complete here
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
GRAVITY = [0 -9.81 0]';

% Calculate zeta_{ij} for (i,j)=(1,2), expressed in body basis
for i=1:NBODIES
    
    % Initialize the fixed parameters of the system
    mass(i) = MASS;
    inertia(:,:,i) = [1 0 0;
                      0 1 0;
                      0 0 INERTIA];
                  
    % Spatial mass matrix written in body basis about the center of mass
    M(:,:,i) = [inertia(:,:,i) zeros(3,3);
                zeros(3,3) mass(i)*eye(3,3)];
    
    length(i) = LENGTH;
    
    % Vectors locating handles 1 and 2 with respect to center of mass
    % Written in body basis. 
    cmtoH1(:,i) = [-length(i)/2 0 0]';
    cmtoH2(:,i) = [length(i)/2 0 0]';
    
    % Written in body basis
    cmtoH1cross(:,:,i) = crossmat(cmtoH1(:,i)); 
                     
    cmtoH2cross(:,:,i) = crossmat(cmtoH2(:,i));
                     
    % Shift matrices for spatial constraint forces on handles H1 and H2
    SfH1(:,:,i) = [eye(3) cmtoH1cross(:,:,i);zeros(3) eye(3)];
    SfH2(:,:,i) = [eye(3) cmtoH2cross(:,:,i);zeros(3) eye(3)];
    
    % handle equations for each body are 
    % A_1 = zeta11 F1 + zeta12 F2 + zeta13
    % A_2 = zeta21 F1 + zeta22 F2 + zeta23
    % phi11, phi12, phi21 and phi22 are constant in body basis and computed
    % just once as a preprocessing step. 
    
    shift_inverse_H1(:,:,i) = (SfH1(:,:,i)')*inv(M(:,:,i));
    shift_inverse_H2(:,:,i) = (SfH2(:,:,i)')*inv(M(:,:,i));
    
    zeta11(:,:,i) = shift_inverse_H1(:,:,i)*SfH1(:,:,i); % (1a)
    zeta12(:,:,i) = shift_inverse_H1(:,:,i)*SfH2(:,:,i); % (1a)
    
    zeta21(:,:,i) = shift_inverse_H2(:,:,i)*SfH1(:,:,i); % (1b)
    zeta22(:,:,i) = shift_inverse_H2(:,:,i)*SfH2(:,:,i); % (1b)      
end

% For a revolute joint about positive z axis
% H is referred to as P^J_k in our paper. 
Hk = [0 0 1 0 0 0]';

% D = D^J_k expressed in body basis
Dk = [1 0 0 0 0;
      0 1 0 0 0;
      0 0 0 0 0;
      0 0 1 0 0;
      0 0 0 1 0;
      0 0 0 0 1];
 
% Everything till this point was preprocessing

% Calculate the transformation matrices and angular velocity.
for i=1:NBODIES    
    
    % Transformation matrix between body and newtonian frame
    n_c_k(:,:,i) = [cos(sum(q(1:i))) -sin(sum(q(1:i))) 0;
                     sin(sum(q(1:i))) cos(sum(q(1:i))) 0;
                     0                0                1];
                     
    N_C_K(:,:,i) = [n_c_k(:,:,i) zeros(3,3);
                    zeros(3,3)  n_c_k(:,:,i)];    
                
    % Angular velocity of each body is trivially obtained here
    W(:,i) = [0 0 sum(u(1:i))]';        
end

% Form the handle equations of motion in Newtonian basis.
% _n is used to denote zeta_{ij} in newtonian basis
for i=1:NP
    for j=1:NB_SYS 
        bid = SID(i,j);
        
        % convert inverse ABI to newtonian basis
        zeta11_n(:,:,bid) = N_C_K(:,:,bid)*zeta11(:,:,bid)*N_C_K(:,:,bid)';
        zeta12_n(:,:,bid) = N_C_K(:,:,bid)*zeta12(:,:,bid)*N_C_K(:,:,bid)';
        zeta21_n(:,:,bid) = N_C_K(:,:,bid)*zeta21(:,:,bid)*N_C_K(:,:,bid)';
        zeta22_n(:,:,bid) = N_C_K(:,:,bid)*zeta22(:,:,bid)*N_C_K(:,:,bid)';
        
        % Written in body basis
        Fa  = [ -cross(W(:,bid),inertia(:,:,bid)*W(:,bid)); % -w x Iw term
                mass(bid)*n_c_k(:,:,bid)'*GRAVITY];% Only external force is        
         
        % written in body basis
        At_H1  = [zeros(3,1);cross(W(:,bid),cross(W(:,bid),cmtoH1(:,bid)))];
        At_H2  = [zeros(3,1);cross(W(:,bid),cross(W(:,bid),cmtoH2(:,bid)))];
        
        % zeta13 and zeta23 written in body basis
        zeta13(:,bid) = (shift_inverse_H1(:,:,bid)*Fa + At_H1); % (1a)
        zeta23(:,bid) = shift_inverse_H2(:,:,bid)*Fa + At_H2; % (1b)
        
        % zeta13 and zeta23 written in Newtonian basis
        zeta13_n(:,bid) = N_C_K(:,:,bid)*zeta13(:,bid);
        zeta23_n(:,bid) = N_C_K(:,:,bid)*zeta23(:,bid);
    end    
end

% Assmelby operation
for i=1:NBODIES-1
        
    pk = ASMID_T(i,1);          % Parent index
    k = ASMID_T(i,2);           % Child index
    pk_k = ASMID_T(i,3);        % Assembled parent-child index
            
    % Assemble k and pk
    X_tilde = inv(Dk'*(zeta11_n(:,:,k) + zeta22_n(:,:,pk))*Dk); % (3e)
    X = Dk*X_tilde*Dk'; % (3d)
    
    Wa(:,:,i) = X*zeta21_n(:,:,pk);               % (3b)
    Wb(:,:,i) = -X*zeta12_n(:,:,k);               % (3b)
    Wc(:,i) = X*(zeta23_n(:,pk)-zeta13_n(:,k));   % (3c)    
    
    % Below are equations (4a)-(4f)
    zeta11_n(:,:,pk_k) = zeta11_n(:,:,pk) - zeta12_n(:,:,pk)*Wa(:,:,i);
    zeta12_n(:,:,pk_k) = -zeta12_n(:,:,pk)*Wb(:,:,i);
    zeta13_n(:,pk_k) = zeta13_n(:,pk) - zeta12_n(:,:,pk)*Wc(:,i);
    
    zeta21_n(:,:,pk_k) = zeta21_n(:,:,k)*Wa(:,:,i);             
    zeta22_n(:,:,pk_k) = zeta22_n(:,:,k) + zeta21_n(:,:,k)*Wb(:,:,i);  
    zeta23_n(:,pk_k) = zeta23_n(:,k) + zeta21_n(:,:,k)*Wc(:,i);          
end

% Solve the root node equations for boundary constraints.
Fh1(:,end) = Dk*((Dk'*zeta11_n(:,:,end)*Dk) \ (-Dk'*zeta13_n(:,end)));
Fh2(:,end) = zeros(6,1);
A1 = zeta11_n(:,:,end)*Fh1(:,end) + zeta13_n(:,end);
udot(1,1) = Hk'*A1;

% disassembly
for i=NBODIES-1:-1:1
    pk = ASMID_T(i,1);          % Parent index
    k = ASMID_T(i,2);           % Child index    
    udot_id = ASMID_T(i,4);     % udot index   
            
    % This is the negative of eq (3a)
    Fc = -Wa(:,:,i)*Fh1(:,i) - Wb(:,:,i)*Fh2(:,i) - Wc(:,i);
    
    % We only ever disassemble assemblies, hence the constraint forces
    % acting on boundary handles of assemblies are the only ones stored
    % resultin the following tow conditional statements
    if(pk>NBODIES)
        id_parent_asm = find(~(ASMID_T(:,3)-pk));
        Fh1(:,id_parent_asm) = Fh1(:,i);
        Fh2(:,id_parent_asm) = Fc;
    end
    
    if(k>NBODIES)
        id_child_asm = find(~(ASMID_T(:,3)-k));
        Fh2(:,id_child_asm) = Fh2(:,i);
        Fh1(:,id_child_asm) = -Fc;
    end
    
    A_pk_2 = zeta21_n(:,:,pk)*Fh1(:,i) + zeta22_n(:,:,pk)*Fc + zeta23_n(:,pk);
    A_k_1 = -zeta11_n(:,:,k)*Fc + zeta12_n(:,:,k)*Fh2(:,i) + zeta13_n(:,k);
    udot(udot_id,1) = Hk'*(A_k_1-A_pk_2);        
end